import json
import magic_pdf.model as model_config
from magic_pdf.pipe.UNIPipe import UNIPipe
from magic_pdf.rw.DiskReaderWriter import DiskReaderWriter
from minio.error import S3Error
import os
import re
import uuid
from docx import Document
from lxml import etree
from loguru import logger
from minio import Minio


class MinIOClient:
    def __init__(self, endpoint, access_key, secret_key, secure=False):
        self.client = Minio(
            endpoint,
            access_key=access_key,
            secret_key=secret_key,
            secure=secure
        )
        self.endpoint = endpoint
        self.secure = secure

    def create_bucket(self, bucket_name):
        try:
            if not self.client.bucket_exists(bucket_name):
                self.client.make_bucket(bucket_name)
            print(f"Bucket '{bucket_name}' is ready.")
        except S3Error as e:
            print(f"Error checking/creating bucket: {e}")

    def set_bucket_policy(self, bucket_name):
        try:
            policy = {
                "Version": "2012-10-17",
                "Statement": [
                    {
                        "Effect": "Allow",
                        "Principal": {"AWS": ["*"]},
                        "Action": ["s3:GetObject"],
                        "Resource": [f"arn:aws:s3:::{bucket_name}/*"]
                    }
                ]
            }
            policy_json = json.dumps(policy)
            self.client.set_bucket_policy(bucket_name, policy_json)
            print(f"Bucket '{bucket_name}' is now publicly accessible.")
        except S3Error as e:
            print(f"Error setting bucket policy: {e}")

    def upload_file(self, bucket_name, file_path, object_name):
        try:
            self.client.fput_object(bucket_name, object_name, file_path)
            print(f"'{object_name}' is successfully uploaded.")
        except S3Error as e:
            print(f"Error uploading file: {e}")
            return None

        protocol = "http" if not self.secure else "https"
        object_url = f"{protocol}://{self.endpoint}/{bucket_name}/{object_name}"
        return object_url


class PDFToMarkdownConverter:
    def __init__(self, output_dir):
        self.output_dir = output_dir
        self.image_pattern = re.compile(r'!\[.*?\]\((.*?)\)')
        model_config.__use_inside_model__ = True

    def convert_pdf_to_markdown(self, pdf_file_path):
        """使用 magic-pdf 将 PDF 转换为 Markdown 并保存到指定目录"""
        try:
            demo_name = os.path.splitext(os.path.basename(pdf_file_path))[0]
            pdf_bytes = open(pdf_file_path, "rb").read()
            model_json = []
            jso_useful_key = {"_pdf_type": "", "model_list": model_json}

            local_image_dir = os.path.join(self.output_dir, 'images')
            image_writer = DiskReaderWriter(local_image_dir)

            # 使用 UNIPipe 解析 PDF 并生成 Markdown
            pipe = UNIPipe(pdf_bytes, jso_useful_key, image_writer)
            pipe.pipe_classify()

            if len(model_json) == 0 and model_config.__use_inside_model__:
                pipe.pipe_analyze()

            pipe.pipe_parse()
            md_content = pipe.pipe_mk_markdown('images', drop_mode="none")

            md_file_path = os.path.join(self.output_dir, f"{demo_name}.md")
            with open(md_file_path, "w", encoding="utf-8") as f:
                f.write(md_content)

            print(f"Markdown file saved as: {md_file_path}")
            return md_file_path, local_image_dir
        except Exception as e:
            logger.exception(e)
            return None, None

    def process_markdown_file(self, md_content, images_dir, minio_client, bucket_name):
        """处理 Markdown 文件中的图像，并将图像上传到 MinIO"""
        image_links = self.image_pattern.findall(md_content)
        if not image_links:
            print("No images found in the markdown file.")
            return md_content

        for image_link in image_links:
            local_image_path = os.path.join(images_dir, os.path.basename(image_link))
            if os.path.exists(local_image_path):
                object_name = f"kb/{os.path.basename(image_link)}"
                public_url = minio_client.upload_file(bucket_name, local_image_path, object_name)

                if public_url:
                    md_content = md_content.replace(image_link, public_url)
            else:
                print(f"Image not found: {local_image_path}")

        return md_content


class WordToMarkdownConverter:
    def __init__(self, output_dir):
        self.output_dir = output_dir
        self.image_pattern = re.compile(r'!\[.*?\]\((.*?)\)')

    def extract_images_and_convert_to_markdown(self, word_file_path):
        """使用 lxml 解析 Word 文档中的图片并生成 Markdown 文件，保持图片位置一致"""
        try:
            demo_name = os.path.splitext(os.path.basename(word_file_path))[0]
            local_image_dir = os.path.join(self.output_dir, 'images')
            os.makedirs(local_image_dir, exist_ok=True)

            # 打开 Word 文档
            doc = Document(word_file_path)
            md_content = ""
            image_count = 0

            # 获取文档的 XML
            xml_content = doc._element
            tree = etree.ElementTree(xml_content)

            # 遍历文档的 XML 结构，保持图片和文本的位置
            for element in tree.iter():
                if element.tag.endswith('p'):
                    # 段落文本
                    paragraph_text = "".join([node.text for node in element.iter() if node.text])
                    md_content += paragraph_text + "\n"

                if element.tag.endswith('blip'):
                    # 提取图片
                    image_count += 1
                    image_filename = f"image{image_count}.png"
                    image_path = os.path.join(local_image_dir, image_filename)

                    # 从图片关系中获取图片数据
                    rId = element.get('{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed')
                    image_part = doc.part.related_parts[rId]

                    # 保存图片
                    with open(image_path, "wb") as f:
                        f.write(image_part.blob)

                    # 插入图片链接到 Markdown 内容
                    md_content += f"![{image_filename}](images/{image_filename})\n"

            # 保存生成的 Markdown 文件
            md_file_path = os.path.join(self.output_dir, f"{demo_name}.md")
            with open(md_file_path, "w", encoding="utf-8") as f:
                f.write(md_content)

            print(f"Markdown file saved as: {md_file_path}")
            print(f"Extracted {image_count} images to {local_image_dir}")
            return md_file_path, local_image_dir

        except Exception as e:
            logger.exception(e)
            return None, None

    def process_markdown_file(self, md_content, images_dir, minio_client, bucket_name):
        """处理 Markdown 文件中的图像，并将图像上传到 MinIO"""
        image_links = self.image_pattern.findall(md_content)

        if not image_links:
            print("No images found in the markdown file.")
            return md_content

        for image_link in image_links:
            # 获取本地图片路径
            local_image_path = os.path.join(images_dir, os.path.basename(image_link))
            if os.path.exists(local_image_path):
                object_name = f"kb/{os.path.basename(image_link)}"
                public_url = minio_client.upload_file(bucket_name, local_image_path, object_name)

                if public_url:
                    # 替换 Markdown 内容中的图片链接为 MinIO 公共 URL
                    md_content = md_content.replace(image_link, public_url)
                    print(f"Replaced image link: {image_link} with {public_url}")
            else:
                print(f"Image not found: {local_image_path}")

        return md_content


class Pipeline:
    def __init__(self, minio_endpoint, minio_access_key, minio_secret_key, bucket_name, output_dir):
        self.minio_client = MinIOClient(minio_endpoint, minio_access_key, minio_secret_key)
        self.bucket_name = bucket_name
        self.output_dir = output_dir
        self.pdf_converter = PDFToMarkdownConverter(output_dir)
        self.word_converter = WordToMarkdownConverter(output_dir)

        # 创建和设置 MinIO 存储桶
        self.minio_client.create_bucket(self.bucket_name)
        self.minio_client.set_bucket_policy(self.bucket_name)

    def run(self, pdf_file_path, original_filename):
        try:
            # 转换 PDF 并生成临时的 Markdown 文件
            temp_md_file_path, images_dir = self.pdf_converter.convert_pdf_to_markdown(pdf_file_path)

            if temp_md_file_path and images_dir:
                # 创建新的 Markdown 文件路径，用于替换临时文件路径
                md_file_path = os.path.join(self.output_dir, f"{original_filename}.md")

                # 读取生成的临时 Markdown 文件内容
                with open(temp_md_file_path, "r", encoding="utf-8") as file:
                    md_content = file.read()

                # 处理 Markdown 中的图像并上传图片到 MinIO
                updated_md_content = self.pdf_converter.process_markdown_file(md_content, images_dir, self.minio_client,
                                                                              self.bucket_name)

                # 将最终内容写入新的 Markdown 文件
                with open(md_file_path, "w", encoding="utf-8") as file:
                    file.write(updated_md_content)

                return md_file_path
            else:
                return None
        finally:
            # 清理临时文件
            if os.path.exists(pdf_file_path):
                os.remove(pdf_file_path)
            if os.path.exists(temp_md_file_path):
                os.remove(temp_md_file_path)

    def process_word(self, word_file_path, original_filename):
        try:
            # 转换 Word 并生成临时的 Markdown 文件
            temp_md_file_path, images_dir = self.word_converter.extract_images_and_convert_to_markdown(word_file_path)

            if temp_md_file_path and images_dir:
                # 创建新的 Markdown 文件路径，用于替换临时文件路径
                md_file_path = os.path.join(self.output_dir, f"{original_filename}.md")

                # 读取生成的临时 Markdown 文件内容
                with open(temp_md_file_path, "r", encoding="utf-8") as file:
                    md_content = file.read()

                # 处理 Markdown 中的图像并上传图片到 MinIO
                updated_md_content = self.word_converter.process_markdown_file(md_content, images_dir,
                                                                               self.minio_client, self.bucket_name)

                # 将最终内容写入新的 Markdown 文件
                with open(md_file_path, "w", encoding="utf-8") as file:
                    file.write(updated_md_content)

                return md_file_path
            else:
                return None
        finally:
            # 清理临时文件
            if os.path.exists(word_file_path):
                os.remove(word_file_path)
            if os.path.exists(temp_md_file_path):
                os.remove(temp_md_file_path)
